#include "matrix_multiplication.h"
#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

typedef float lf;

void __global__ MatMatMul(lf *At, lf *B, lf *C, size_t m, size_t n, size_t p) {
    extern lf __shared__ shared[];
    lf res = 0, *Ats = shared, *Bs = shared + blockDim.x * blockDim.y;
    size_t r = blockIdx.y * blockDim.x + threadIdx.y,
            c = blockIdx.x * blockDim.x + threadIdx.x;
    for (size_t t = 0; t < n; t += blockDim.x) {
        __syncthreads();
        Ats[threadIdx.y * blockDim.x + threadIdx.x] = r < m && t + threadIdx.x < n ? At[(t + threadIdx.x) * m + r] : 0;
        Bs[threadIdx.x * blockDim.y + threadIdx.y] = c < p && t + threadIdx.y < n ? B[(t + threadIdx.y) * p + c] : 0;
        __syncthreads();
        for (size_t i = 0; i < blockDim.x; ++i)
            res += Ats[i * blockDim.y + threadIdx.y] * Bs[i * blockDim.x + threadIdx.x];
    }
    if (r < m && c < p)
        C[r * p + c] = res;
}

void matmulexpr() {
    const size_t
            nRow = 1 << 12,
            nCol = 1 << 12;
    lf *h_A = (lf *) malloc(sizeof(lf) * nRow * nCol),
            *h_At = (lf *) malloc(sizeof(lf) * nRow * nCol),
            *h_B = (lf *) malloc(sizeof(lf) * nCol * nRow),
            *d_A,
            *d_At,
            *d_B,
            *d_C;
    for (size_t i = 0; i < nRow; ++i)
        for (size_t j = 0; j < nCol; ++j) {
            h_A[i * nCol + j] = (i - 0.1 * j + 1) / (i + j + 1);
            h_At[j * nRow + i] = (i - 0.1 * j + 1) / (i + j + 1);
            h_B[i * nCol + j] = (j - 0.2 * i + 1) * (i + j + 1) / (i * i + j * j + 1);
        }

    cudaMalloc((lf **) &d_A, sizeof(lf) * nRow * nCol);
    cudaMalloc((lf **) &d_At, sizeof(lf) * nRow * nCol);
    cudaMalloc((lf **) &d_B, sizeof(lf) * nCol * nRow);
    cudaMalloc((lf **) &d_C, sizeof(lf) * nRow * nRow);

    cudaMemcpy(d_A, h_A, sizeof(lf) * nRow * nCol, cudaMemcpyHostToDevice);
    cudaMemcpy(d_At, h_At, sizeof(lf) * nRow * nCol, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, sizeof(lf) * nCol, cudaMemcpyHostToDevice);

    for (int tile_width = 32; tile_width; tile_width >>= 1) {
        cudaEvent_t beg, end;
        cudaEventCreate(&beg);
        cudaEventCreate(&end);
        cudaEventRecord(beg, 0);

        dim3 grids((nRow - 1) / tile_width + 1, (nRow - 1) / tile_width + 1),
                blocks(tile_width, tile_width);
        MatMatMul<<<grids, blocks, sizeof(lf) * 2 * blocks.x * blocks.y>>>(d_At, d_B, d_C, nRow, nCol, nRow);

        cudaDeviceSynchronize();
        cudaEventRecord(end, 0);
        cudaEventSynchronize(beg);
        cudaEventSynchronize(end);
        lf elapsed_time;
        cudaEventElapsedTime(&elapsed_time, beg, end);
        printf("(%d,%d)grids, (%d,%d)blocks, %fms elapsed.\n", grids.x, grids.y, blocks.x, blocks.y, elapsed_time);
    }

    cudaFree(d_A);
    cudaFree(d_At);
    cudaFree(d_B);
    cudaFree(d_C);
    free(h_A);
    free(h_At);
    free(h_B);
}
